//! # 异步任务池实现
//! 
//! 提供高效的异步任务调度和执行机制，支持任务队列管理和工作线程池。
//! 通过合理的任务分发和负载均衡提升系统并发处理能力。

use super::traits::{AsyncTaskPool, TaskStats};
use rustcloud_core::{ServiceResult, ServiceError};
use std::sync::Arc;
use tokio::sync::{RwLock, mpsc, Semaphore};
use std::future::Future;

/// 异步任务池实现
pub struct AsyncTaskPoolImpl {
    /// 任务发送器
    task_sender: mpsc::UnboundedSender<BoxedTask>,
    /// 任务接收器
    task_receiver: Arc<mpsc::UnboundedReceiver<BoxedTask>>,
    /// 工作线程数
    worker_count: usize,
    /// 任务统计
    stats: Arc<RwLock<TaskStats>>,
    /// 关闭信号
    shutdown_semaphore: Arc<Semaphore>,
}

type BoxedTask = Box<dyn FnOnce() + Send + 'static>;

impl AsyncTaskPoolImpl {
    /// 创建异步任务池
    pub fn new(worker_count: usize) -> Self {
        let (task_sender, task_receiver) = mpsc::unbounded_channel::<BoxedTask>();
        let stats = Arc::new(RwLock::new(TaskStats::default()));
        let shutdown_semaphore = Arc::new(Semaphore::new(0));

        let pool = Self {
            task_sender,
            task_receiver: Arc::new(task_receiver),
            worker_count,
            stats: stats.clone(),
            shutdown_semaphore: shutdown_semaphore.clone(),
        };

        // 启动工作线程
        pool.start_workers(stats, shutdown_semaphore);

        pool
    }

    /// 启动工作线程
    fn start_workers(&self, stats: Arc<RwLock<TaskStats>>, shutdown_semaphore: Arc<Semaphore>) {
        for i in 0..self.worker_count {
            let mut receiver = self.task_receiver.clone();
            let stats_clone = stats.clone();
            let shutdown_clone = shutdown_semaphore.clone();
            
            tokio::spawn(async move {
                tracing::debug!("异步任务池工作线程 {} 启动", i);
                
                loop {
                    tokio::select! {
                        // 接收任务
                        task = receiver.recv() => {
                            if let Some(task) = task {
                                // 执行任务
                                task();
                                
                                // 更新统计
                                let mut stats = stats_clone.write().await;
                                stats.completed_tasks += 1;
                            } else {
                                // 通道已关闭
                                break;
                            }
                        }
                        // 检查关闭信号
                        _ = shutdown_clone.acquire() => {
                            // 收到关闭信号
                            tracing::debug!("异步任务池工作线程 {} 收到关闭信号", i);
                            break;
                        }
                    }
                }
                
                tracing::debug!("异步任务池工作线程 {} 退出", i);
            });
        }

        // 更新统计中的工作线程数
        tokio::spawn(async move {
            let mut stats = stats.write().await;
            stats.worker_count = worker_count;
        });
    }
}

#[async_trait::async_trait]
impl AsyncTaskPool for AsyncTaskPoolImpl {
    async fn submit<F>(&self, task: F) -> ServiceResult<()>
    where
        F: FnOnce() + Send + 'static,
    {
        // 更新统计
        {
            let mut stats = self.stats.write().await;
            stats.submitted_tasks += 1;
            stats.queue_length += 1;
        }

        // 发送任务
        self.task_sender.send(Box::new(task))
            .map_err(|_| ServiceError::InternalError("任务队列已关闭".to_string()))?;

        Ok(())
    }

    async fn submit_async<F, T>(&self, task: F) -> ServiceResult<tokio::task::JoinHandle<T>>
    where
        F: Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        // 更新统计
        {
            let mut stats = self.stats.write().await;
            stats.submitted_tasks += 1;
            stats.queue_length += 1;
        }

        // 使用tokio::spawn执行异步任务
        let handle = tokio::spawn(async move {
            task.await
        });

        Ok(handle)
    }

    async fn get_stats(&self) -> TaskStats {
        self.stats.read().await.clone()
    }

    async fn shutdown(&self) -> ServiceResult<()> {
        // 发送关闭信号给所有工作线程
        for _ in 0..self.worker_count {
            let _ = self.shutdown_semaphore.try_acquire();
        }

        // 关闭任务发送通道
        drop(self.task_sender.clone());

        // 等待一段时间让工作线程处理完剩余任务
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

        tracing::info!("异步任务池已关闭");
        Ok(())
    }

    async fn force_shutdown(&self) -> ServiceResult<()> {
        // 直接关闭任务发送通道
        drop(self.task_sender.clone());

        tracing::info!("异步任务池已强制关闭");
        Ok(())
    }
}

/// 异步任务池构建器
pub struct AsyncTaskPoolBuilder {
    worker_count: usize,
}

impl AsyncTaskPoolBuilder {
    /// 创建新的构建器
    pub fn new() -> Self {
        Self {
            worker_count: num_cpus::get(),
        }
    }

    /// 设置工作线程数
    pub fn worker_count(mut self, count: usize) -> Self {
        self.worker_count = count;
        self
    }

    /// 构建异步任务池
    pub fn build(self) -> AsyncTaskPoolImpl {
        AsyncTaskPoolImpl::new(self.worker_count)
    }
}

impl Default for AsyncTaskPoolBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// 便捷函数：创建默认配置的异步任务池
pub fn create_default_async_pool() -> AsyncTaskPoolImpl {
    AsyncTaskPoolBuilder::new().build()
}

/// 便捷函数：创建指定工作线程数的异步任务池
pub fn create_async_pool_with_workers(worker_count: usize) -> AsyncTaskPoolImpl {
    AsyncTaskPoolBuilder::new().worker_count(worker_count).build()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::time::Duration;

    #[tokio::test]
    async fn test_async_task_pool_basic_operations() {
        let pool = AsyncTaskPoolImpl::new(4);

        let counter = Arc::new(AtomicU32::new(0));
        let counter_clone = counter.clone();

        // 提交任务
        pool.submit(move || {
            counter_clone.fetch_add(1, Ordering::Relaxed);
        }).await.unwrap();

        // 等待任务执行
        tokio::time::sleep(Duration::from_millis(50)).await;

        assert_eq!(counter.load(Ordering::Relaxed), 1);

        // 验证统计信息
        let stats = pool.get_stats().await;
        assert_eq!(stats.submitted_tasks, 1);
        assert_eq!(stats.completed_tasks, 1);
        assert_eq!(stats.worker_count, 4);

        // 关闭任务池
        pool.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_async_task_pool_with_async_tasks() {
        let pool = AsyncTaskPoolImpl::new(2);

        // 提交异步任务
        let handle = pool.submit_async(async {
            tokio::time::sleep(Duration::from_millis(10)).await;
            42
        }).await.unwrap();

        // 等待任务完成并获取结果
        let result = handle.await.unwrap();
        assert_eq!(result, 42);

        // 验证统计信息
        let stats = pool.get_stats().await;
        assert_eq!(stats.submitted_tasks, 1);
        assert_eq!(stats.completed_tasks, 1);

        pool.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_async_task_pool_multiple_tasks() {
        let pool = AsyncTaskPoolImpl::new(3);
        let counter = Arc::new(AtomicU32::new(0));

        // 提交多个任务
        for i in 0..10 {
            let counter_clone = counter.clone();
            pool.submit(move || {
                // 模拟一些工作
                std::thread::sleep(Duration::from_millis(5));
                counter_clone.fetch_add(i, Ordering::Relaxed);
            }).await.unwrap();
        }

        // 等待所有任务完成
        tokio::time::sleep(Duration::from_millis(100)).await;

        // 验证所有任务都已执行
        let sum: u32 = (0..10).sum();
        assert_eq!(counter.load(Ordering::Relaxed), sum);

        // 验证统计信息
        let stats = pool.get_stats().await;
        assert_eq!(stats.submitted_tasks, 10);
        assert_eq!(stats.completed_tasks, 10);

        pool.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_async_task_pool_builder() {
        let pool = AsyncTaskPoolBuilder::new()
            .worker_count(6)
            .build();

        let stats = pool.get_stats().await;
        assert_eq!(stats.worker_count, 6);

        pool.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_async_task_pool_concurrent_submission() {
        let pool = AsyncTaskPoolImpl::new(4);
        let counter = Arc::new(AtomicU32::new(0));

        // 并发提交任务
        let mut handles = vec![];
        for i in 0..20 {
            let counter_clone = counter.clone();
            let pool_clone = pool.clone();
            let handle = tokio::spawn(async move {
                pool_clone.submit(move || {
                    std::thread::sleep(Duration::from_millis(2));
                    counter_clone.fetch_add(i, Ordering::Relaxed);
                }).await.unwrap();
            });
            handles.push(handle);
        }

        // 等待所有提交完成
        for handle in handles {
            handle.await.unwrap();
        }

        // 等待所有任务执行完成
        tokio::time::sleep(Duration::from_millis(100)).await;

        // 验证统计信息
        let stats = pool.get_stats().await;
        assert_eq!(stats.submitted_tasks, 20);
        assert!(stats.completed_tasks >= 15); // 可能还有一些在执行中

        pool.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_force_shutdown() {
        let pool = AsyncTaskPoolImpl::new(2);

        // 提交一个长时间运行的任务
        pool.submit(|| {
            std::thread::sleep(Duration::from_secs(5));
        }).await.unwrap();

        // 强制关闭应该立即返回
        let start = std::time::Instant::now();
        pool.force_shutdown().await.unwrap();
        let duration = start.elapsed();

        assert!(duration < Duration::from_secs(1));
    }
}